"""
plot uncertainty lines, band
calculate variation
"""
# -*- coding : utf-8 -*-
import os
import copy
import matplotlib.pyplot as plt
import numpy as np
import gauge
import wrf_output


def wipeOff_max_min(arr):
    """calculate mean except for min and max"""
    max_arr = np.max(arr)
    min_arr = np.min(arr)
    if max_arr != min_arr:
        arr[np.logical_and(arr == max_arr, arr == min_arr)] = np.nan
    return arr


def decumulation(seq):
    """de-accumulation"""
    result = [seq[i + 1] - seq[i] for i in range(len(seq) - 1)]
    result.insert(0, seq[0])
    return result


def get_sim_accum(dat_path, grid_spacing, domain, start_time, mps, cus, pbls):
    """get simulations, accumulation according to different combinations of mp, cu and pbl"""
    gauge_loc = ''
    if grid_spacing == '139':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge/Gauge_row_col_1.csv'
    elif grid_spacing == '51545':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge/Gauge_row_col_5.csv'
    elif grid_spacing == '103090':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge/Gauge_row_col_10.csv'

    accumulation = []
    simulations = []
    labels = []
    for i in mps:
        for j in cus:
            for k in pbls:
                mp_cu_pbl = i + j + k
                if os.path.exists(dat_path + '/' + mp_cu_pbl):
                    labels.append(i + '-' + j + '-' + k)
                    input_path = dat_path + '/' + mp_cu_pbl + \
                        '/wrfout_d' + domain + '_' + start_time
                    all_gauges = wrf_output.GetAll(input_path, gauge_loc)

                    # wipeOff max and min value
                    # all_gauges_filter = np.array(
                    #     list(map(wipeOff_max_min, all_gauges)))
                    # average = np.mean(all_gauges_filter, axis=1)

                    # add accumulation
                    accumulation_average = np.mean(all_gauges, axis=1)
                    accumulation.append(copy.deepcopy(accumulation_average))
                    # add simulation
                    hourly_simulation = decumulation(all_gauges)
                    hourly_average = np.mean(hourly_simulation, axis=1)
                    simulations.append(hourly_average)

    return simulations, accumulation, labels


def plot_lines_bands(dat_path, domain, start_time, event, simulations, gauges_mean, accumulation, gauges_accum, labels):
    """plot lines, band"""
    mps = ['2', '6', '7']
    cus = ['0', '1', '2', '7']
    pbls = ['2', '6', '9']
    grid_spacing = dat_path.split('/')[-1]

    # hourly band and accumulative band
    # hourly_band = [np.min(simulations, axis=0), np.max(simulations, axis=0)]
    # accum_band = [np.min(accumulation, axis=0), np.max(accumulation, axis=0)]

    sim_min = np.min(simulations, axis=0)
    sim_max = np.max(simulations, axis=0)
    sim_5 = sim_min + (sim_max - sim_min) * 0.05
    sim_95 = sim_min + (sim_max - sim_min) * 0.95
    hourly_band = np.array([sim_5, sim_95])

    accum_min = np.min(accumulation, axis=0)
    accum_max = np.max(accumulation, axis=0)
    accum_5 = accum_min + (accum_max - accum_min) * 0.05
    accum_95 = accum_min + (accum_max - accum_min) * 0.95
    accum_band = [accum_5, accum_95]

    # calculate deviation
    deviations = []
    for _, result in enumerate(simulations):
        deviation = result - gauges_mean
        deviations.append(np.mean(np.abs(deviation)))
    min_index = np.argmin(deviations)
    min_label = labels[min_index]
    # plot
    plt.style.use('fivethirtyeight')
    labels = ('integrated ensemble', 'observations',
              'ensemble average', min_label)
    indexes = gauges_mean.index
    time_label = [str(hour) for hour in indexes.hour]
    x_series = np.arange(len(time_label))
    hourly_crest = 0
    cum_crest = 0
    if event == "R1":
        hourly_crest = 6
        cum_crest = 60
    elif event == "R2":
        hourly_crest = 8
        cum_crest = 110
    elif event == "R3":
        hourly_crest = 4
        cum_crest = 50
    elif event == "R4":
        hourly_crest = 8
        cum_crest = 80
    elif event == "R5":
        hourly_crest = 4
        cum_crest = 70
    elif event == "R6":
        hourly_crest = 6
        cum_crest = 50
    elif event == "R7":
        hourly_crest = 5
        cum_crest = 40
    elif event == "R8":
        hourly_crest = 4
        cum_crest = 40

    _, ax = plt.subplots(nrows=3, ncols=1, sharex=True, figsize=(15, 15))
    # plot lines from simulations
    for _, line in enumerate(simulations):
        ax[0].plot(x_series, line, lw=1, alpha=1,
                   color="#539caf", linestyle=(0, (1, 5)))
    # plot band
    # ax[0].fill_between(x_series, hourly_band[0], hourly_band[1], label=labels[0],
    #                    color='#539caf', alpha=0.4)
    ax[0].fill_between(x_series, hourly_band[0], hourly_band[1],
                       color='#539caf', alpha=0.4)
    # ax[0].plot(x_series, np.mean(hourly_band, axis=0), color='#539caf',
    #            lw=2, alpha=1, label=labels[2])
    ax[0].plot(x_series, np.mean(hourly_band, axis=0), color='#539caf',
               lw=2, alpha=1)
    # plot the line with min deviation
    ax[0].plot(x_series, simulations[min_index], color='#7663b0',
               label=labels[3], lw=2, alpha=1)
    # plot gauges mean
    # ax[0].plot(x_series, gauges_mean, color='#FF7F0E', label=labels[1],
    #            lw=2, alpha=1)
    ax[0].plot(x_series, gauges_mean, color='#FF7F0E',
               lw=2, alpha=1)

    ax[0].yaxis.set_tick_params(labelsize=20)
    ax[0].set_title(event + " hourly rainfall", fontsize=24)
    # ax[0].set_yticks(np.arange(0, max(max(band[1]), max(gauges_mean)) + 1, 2))
    ax[0].set_yticks(np.arange(0, hourly_crest + 2, 2))
    ax[0].legend(loc='upper left', fontsize='large')

    # plot mp, cu and pbl ensembles

    min_labels = min_label.split('-')
    min_mp = [min_labels[0], ]
    min_cu = [min_labels[1], ]
    min_pbl = [min_labels[2], ]
    # calculate mp ensemble
    mp_simulations, _, _ = get_sim_accum(
        dat_path, grid_spacing, domain, start_time, mps, min_cu, min_pbl)
    mp_simulations = np.asarray(mp_simulations)
    # mp_band = [np.min(mp_simulations, axis=0), np.max(mp_simulations, axis=0)]
    mp_min = np.min(mp_simulations, axis=0)
    mp_max = np.max(mp_simulations, axis=0)
    mp_band = [mp_min + (mp_max - mp_min) * 0.05,
               mp_min + (mp_max - mp_min) * 0.95]
    # calculate cu ensemble
    cu_simulations, _, _ = get_sim_accum(
        dat_path, grid_spacing, domain, start_time, min_mp, cus, min_pbl)
    cu_simulations = np.asarray(cu_simulations)
    # cu_band = [np.min(cu_simulations, axis=0), np.max(cu_simulations, axis=0)]
    cu_min = np.min(cu_simulations, axis=0)
    cu_max = np.max(cu_simulations, axis=0)
    cu_band = [cu_min + (cu_max - cu_min) * 0.05,
               cu_min + (cu_max - cu_min) * 0.95]
    # calculate pbl ensemble
    pbl_simulations, _, _ = get_sim_accum(
        dat_path, grid_spacing, domain, start_time, min_mp, min_cu, pbls)
    pbl_simulations = np.asarray(pbl_simulations)
    # pbl_band = [np.min(pbl_simulations, axis=0),
    #             np.max(pbl_simulations, axis=0)]
    pbl_min = np.min(pbl_simulations, axis=0)
    pbl_max = np.max(pbl_simulations, axis=0)
    pbl_band = [pbl_min + (pbl_max - pbl_min) * 0.05,
                pbl_min + (pbl_max - pbl_min) * 0.95]
    ensemble_labels = ('integrated ensemble', 'mp-ensemble',
                       'cu-ensemble', 'pbl-ensemble')
    ax[1].fill_between(x_series, hourly_band[0], hourly_band[1],
                       color='#539caf', label=ensemble_labels[0], alpha=0.2)
    ax[1].fill_between(x_series, mp_band[0], mp_band[1],
                       color='#1b9e77', label=ensemble_labels[1], alpha=0.4)
    ax[1].fill_between(x_series, cu_band[0], cu_band[1],
                       color='#d95f02', label=ensemble_labels[2], alpha=0.6)
    ax[1].fill_between(x_series, pbl_band[0], pbl_band[1],
                       color='#e7298a', label=ensemble_labels[3], alpha=0.5)
    ax[1].set_ylabel('precipitation (mm)', fontsize=22)
    ax[1].set_yticks(np.arange(0, hourly_crest + 2, 2))
    ax[1].yaxis.set_tick_params(labelsize=20)
    ax[1].set_title("mp, cu and pbl ensembles", fontsize=24)
    # ax[1].legend(loc='upper left', fontsize='large')

    # plot accumulated precipitation
    ax[2].fill_between(x_series, accum_band[0], accum_band[1],
                       color='#539caf', alpha=0.4, label=labels[0])
    ax[2].plot(x_series, np.mean(accum_band, axis=0), color='#539caf',
               lw=2, alpha=1, label=labels[2])
    ax[2].plot(x_series, gauges_accum, color='#FF7F0E',
               lw=2, alpha=1, label=labels[1])
    ax[2].set_xlabel('time (h)', fontsize=22)
    ax[2].xaxis.set_tick_params(labelsize=20)
    ax[2].yaxis.set_tick_params(labelsize=20)
    ax[2].set_yticks(np.arange(0, cum_crest + 20, 20))
    ax[2].set_title("accumulative rainfall", fontsize=24)

    plt.tight_layout()
    plt.show()


def plot_var(dat_path, domain, start_time, end_time, event):
    """plot rainfall variation"""
    mps = ['2', '6', '7']
    cus = ['0', '1', '2', '7']
    pbls = ['2', '6', '9']
    split_path = dat_path.split('/')
    grid_spacing = split_path[-1]

    simulations, accumulation, labels = get_sim_accum(
        dat_path, grid_spacing, domain, start_time, mps, cus, pbls)
    # get simulations
    simulations = np.asarray(simulations)
    # get the mean gauge value
    gauges_mean = gauge.Get50GaugeMean(start_time, end_time)

    # get accumulation
    accumulation = np.asarray(accumulation)
    # get the accumulated gauge value
    gauges_accum = gauge.Get50GaugeMeanAccum(start_time, end_time)

    # plot lines and band
    plot_lines_bands(dat_path, domain, start_time, event,
                     simulations, gauges_mean,
                     accumulation, gauges_accum,
                     labels)


def main():
    """main function"""
    # R1
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/01/1700/103090'
    # domain = r'03'
    # start_time = r'2008-01-17_00_00_00'
    # end_time = r'2008-01-19_12_00_00'
    # event = 'R1'

    # R2
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/01/1912/139'
    # domain = r'03'
    # start_time = r'2008-01-19_12_00_00'
    # end_time = r'2008-01-22_00_00_00'
    # event = 'R2'

    # R3
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/08/1718/139'
    # domain = r'03'
    # start_time = r'2008-08-17_18_00_00'
    # end_time = r'2008-08-20_00_00_00'
    # event = 'R3'

    # R4
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/09/0500/139'
    # domain = r'03'
    # start_time = r'2008-09-05_00_00_00'
    # end_time = r'2008-09-07_00_00_00'
    # event = 'R4'

    # R5
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/09/2900/139'
    # domain = r'03'
    # start_time = r'2008-09-29_00_00_00'
    # end_time = r'2008-10-02_06_00_00'
    # event = 'R5'

    # R6
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/10/2506/139'
    # domain = r'03'
    # start_time = r'2008-10-25_06_00_00'
    # end_time = r'2008-10-26_06_00_00'
    # event = 'R6'

    # R7
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/11/0900/139'
    # domain = r'03'
    # start_time = r'2008-11-09_00_00_00'
    # end_time = r'2008-11-10_06_00_00'
    # event = 'R7'

    # R8
    dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/12/0400/139'
    domain = r'03'
    start_time = r'2008-12-04_00_00_00'
    end_time = r'2008-12-06_00_00_00'
    event = 'R8'

    plot_var(dat_path, domain, start_time, end_time, event)


if __name__ == '__main__':
    main()
